import time
import numpy as np
from agent import SarsaAgent
from environment import Environment


def run_episode(env, agent, render=False):
    total_steps = 0  # 记录每个episode走了多少step
    total_reward = 0

    obs = env.reset()  # 重置环境, 重新开一局（即开始新的一个episode）
    action = agent.sample(obs)  # 根据算法选择一个动作

    while True:
        next_obs, reward, done, _ = env.step(action)  # 与环境进行一个交互
        next_action = agent.sample(next_obs)  # 根据算法选择一个动作
        # 训练 Sarsa 算法
        agent.learn(obs, action, reward, next_obs, next_action, done)

        action = next_action
        obs = next_obs  # 存储上一个观察值
        total_reward += reward
        total_steps += 1  # 计算step数
        if render:
            env.render()  # 渲染新的一帧图形
        if done:
            break
    return total_reward, total_steps


def episode_test(env, agent):
    total_reward = 0
    obs = env.reset()
    while True:
        action = agent.predict(obs)  # greedy
        next_obs, reward, done, _ = env.step(action)
        total_reward += reward
        obs = next_obs
        time.sleep(0.5)
        env.render()
        if done:
            print('test reward = %.1f' % total_reward)
            break

def print_Q_table(Q_table, env_w):
    for i in range(len(Q_table)):
        max_action = np.argmax(Q_table[i])
        if i % env_w == 0:
            print("")
        if max_action == 0:
            print("^", end="")
        elif max_action == 1:
            print("v", end="")
        elif max_action == 2:
            print("<", end="")
        elif max_action == 3:
            print(">", end="")

def main():
    # env = gym.make("CliffWalking-v0")  # 0 up, 1 right, 2 down, 3 left
    # env = CliffWalkingWapper(env)
    init_env = [
        '            ',
        '            ',
        '            ',
        'AXXXXXXXXXXV'
    ]
    # 转换为矩阵形式
    init_env_mat = [list(line) for line in init_env]

    env = Environment(init_env_mat, agent='A', road=' ', trap='X', target='V')

    # 观测状态维度
    observation_dimension = env.observation_dimension
    # 动作维度
    action_dimension = 4

    agent = SarsaAgent(
        obs_n=observation_dimension,
        act_n=action_dimension,
        learning_rate=0.1,
        gamma=0.9,
        e_greed=0.1)

    is_render = False
    for episode in range(500):
        ep_reward, ep_steps = run_episode(env, agent, is_render)
        print('Episode %s: steps = %s , reward = %.1f' % (episode, ep_steps,
                                                          ep_reward))

        # 每隔20个episode渲染一下看看效果
        if episode % 20 == 0:
            is_render = True
        else:
            is_render = False
    # 训练结束，查看算法效果
    episode_test(env, agent)
    # 打印Q表格（可视化）
    print_Q_table(agent.Q, env.weight)


if __name__ == "__main__":
    main()
